from flask import Flask, request, jsonify
from airflow.api.client.local_client import Client
import time

app = Flask(__name__)
client = Client(None)

@app.route('/execute_sql', methods=['POST'])
def trigger_sql_execution():
    try:
        # 获取请求中的 SQL
        data = request.get_json()
        sql = data.get('sql')
        
        if not sql:
            return jsonify({'error': 'No SQL provided'}), 400
        
        # 触发 DAG 运行
        dag_run = client.trigger_dag(
            dag_id='sql_executor',
            conf={'sql': sql}
        )
        
        # 等待 DAG 执行完成
        while True:
            dag_run_status = client.get_dag_run_status(
                dag_id='sql_executor',
                run_id=dag_run.run_id
            )
            if dag_run_status in ['success', 'failed']:
                break
            time.sleep(1)
        
        if dag_run_status == 'failed':
            return jsonify({'error': 'DAG execution failed'}), 500
        
        # 获取执行结果
        result = client.get_task_xcom(
            dag_id='sql_executor',
            task_id='execute_sql',
            run_id=dag_run.run_id,
            key='sql_results'
        )
        
        return jsonify({'results': result})
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)